#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：BroadviewCoss
@File    ：run_task_and_send_alert_msg.py
@IDE     ：PyCharm
@Author  ：kdly
@Date    ：2025/4/15 下午1:28
'''
import math
import time
import logging
import requests

from fake_useragent import UserAgent  # 用于生成随机User-Agent
from django.utils import timezone  # Django时区工具

import utils.jobs as job
from utils.jobs import init_scheduler
from hosts_mgr.models import HostList
from monitor_alert_center.models import MonitorTask, AlertContactGroup, AlertContactPerson, AlertLogs
from utils.encrypt import Password  # 密码加密解密工具
from utils.send_msg_util.send_dingtalk_msg import send_dingtalk_message
from utils.send_msg_util.send_email_utils import send_email
from utils.send_msg_util.send_feishu_msg import send_feishu_message
from utils.send_msg_util.send_wxwork_msg import send_wxwork_message
from utils.ssh_client import LocalSSHClient  # 自定义SSH客户端

# 配置日志
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # 设置日志级别为INFO


class RunTaskAndSendAlertMsg:
    """执行定时任务和发送告警信息的主类"""

    def start(self):
        """启动定时任务调度器"""
        from apscheduler.schedulers.blocking import BlockingScheduler

        self.sched = BlockingScheduler()  # 创建阻塞式调度器
        # 初始化调度器并加载数据库中的任务
        init_scheduler()
        self.sched.start()  # 启动调度器

    def disabled_job(self, types, task_id):
        """
        根据任务类型删除定时任务
        :param types: 监控类型(1：站点检测，2：端口检测，3：进程检测，4：自定义脚本)
        :param task_id: 要删除的任务ID
        """
        if types == 1 or types == 2:
            job.delete_job(f"{task_id}__run_check_port_or_website_wrapper")
        elif types == 3 or types == 4:
            job.delete_job(f"{task_id}__run_cmd_wrapper")

    def run_task(self, type, task_id, mon_obj, mon_freq, host_id, timeout=0):
        """
        创建并运行定时任务
        :param type: 监控类型(1：站点检测，2：端口检测，3：进程检测，4：自定义脚本)
        :param task_id: 任务ID
        :param mon_obj: 监控对象(URL/端口/进程名/脚本)
        :param mon_freq: 监控频率(秒)
        :param host_id: 主机ID(仅类型3和4需要)
        :param timeout: 请求超时时间(毫秒，仅类型1和2需要)
        :return: 是否成功
        """
        try:
            if type == 1 or type == 2:
                # 调度站点/端口检测任务
                job.schedule_job(
                    RunTaskAndSendAlertMsg._run_check_port_or_website_wrapper,
                    mon_freq,
                    task_id,
                    args=(task_id, mon_obj, timeout)
                )
            elif type == 3 or type == 4:
                # 获取主机信息用于进程/脚本检测
                host_obj = HostList.objects.filter(id=host_id).first()
                # 解密密码
                encryption_pass = Password()
                password = encryption_pass.decrypt(host_obj.password)

                # 调度进程/脚本检测任务
                job.schedule_job(
                    RunTaskAndSendAlertMsg._run_cmd_wrapper,
                    mon_freq,
                    task_id,
                    args=(
                        task_id,
                        mon_obj,
                        host_obj.connect_ip,
                        host_obj.port,
                        host_obj.user,
                        password
                    )
                )
        except Exception as e:
            logger.error(e)
            return False
        return True

    @staticmethod
    def _run_check_port_or_website_wrapper(task_id, website_name, timeout=0):
        """
        包装站点/端口检测方法用于调度
        :param task_id: 任务ID
        :param website_name: 网站URL或端口
        :param timeout: 超时时间(毫秒)
        """
        try:
            instance = RunTaskAndSendAlertMsg()
            instance._RunTaskAndSendAlertMsg__run_check_port_or_website(task_id, website_name, timeout)
        except Exception as e:
            logger.error(e)
            raise

    @staticmethod
    def _run_cmd_wrapper(task_id, mon_obj, ip_addr, port, username, password):
        """
        包装命令执行方法用于调度
        :param task_id: 任务ID
        :param mon_obj: 监控对象(进程名或脚本)
        :param ip_addr: 主机IP
        :param port: SSH端口
        :param username: SSH用户名
        :param password: SSH密码
        """
        try:
            instance = RunTaskAndSendAlertMsg()
            instance._RunTaskAndSendAlertMsg__run_cmd(
                task_id=task_id,
                mon_obj=mon_obj,
                ip_addr=ip_addr,
                port=int(port),
                username=username,
                password=password
            )
        except Exception as e:
            logger.error(e)
            raise

    def __run_check_port_or_website(self, task_id, website_name, timeout=0):
        """
        执行站点/端口检测
        :param task_id: 任务ID
        :param website_name: 网站URL或端口
        :param timeout: 超时时间(毫秒)
        """
        mon_task_obj = MonitorTask.objects.filter(id=task_id)
        try:
            ua = UserAgent()
            headers = {'User-Agent': ua.random}  # 随机User-Agent
            if timeout:
                timeout = int(timeout) / 1000  # 转换为秒
                res = requests.get(website_name, headers=headers, timeout=timeout)
            else:
                res = requests.get(website_name, headers=headers)
            logger.info(f'执行计划任务：{task_id}__run_check_port_or_website')

            # 根据状态码判断是否成功
            if 200 <= int(res.status_code) < 500:
                self.__set_fail_num_and_send_success_msg(
                    mon_task_obj,
                    error_msg=f"请求返回状态码：{res.status_code}",
                    task_id=task_id
                )
            else:
                self.__set_fail_num_and_send_fail_msg(
                    mon_task_obj,
                    error_msg=f"请求返回状态码：{res.status_code}",
                    task_id=task_id
                )
        except Exception as e:
            self.__set_fail_num_and_send_fail_msg(mon_task_obj, error_msg=str(e), task_id=task_id)
            logger.error(e)
            raise

    def __run_cmd(self, task_id, mon_obj='', ip_addr='', port=0, username='', password=''):
        """
        执行进程检测或自定义脚本
        :param task_id: 任务ID
        :param mon_obj: 进程名或脚本命令
        :param ip_addr: 主机IP
        :param port: SSH端口
        :param username: SSH用户名
        :param password: SSH密码
        """
        try:
            mon_task_obj = MonitorTask.objects.filter(id=task_id)
            client = LocalSSHClient(ip_addr=ip_addr, port=port, username=username, password=password)

            # 根据任务类型执行不同命令
            if mon_task_obj.get().type == 3:
                # 进程检测命令
                code = client.execute_command(f'ps -ef|grep -v grep|grep {mon_obj} &>/dev/null', 1)
            else:
                # 自定义脚本命令
                code = client.execute_command(f"""{mon_obj}""", 1)

            # 根据返回码判断是否成功
            if int(code) == 0:
                self.__set_fail_num_and_send_success_msg(
                    mon_task_obj,
                    error_msg="进程检测/监控脚本执行成功",
                    task_id=task_id
                )
            else:
                self.__set_fail_num_and_send_fail_msg(
                    mon_task_obj,
                    error_msg="进程检测/监控脚本执行失败",
                    task_id=task_id
                )
        except Exception as e:
            logger.error(e)
            raise

    def __set_fail_num_and_send_success_msg(self, mon_task_obj, error_msg='', task_id=0):
        """
        处理任务成功的情况
        :param mon_task_obj: MonitorTask查询集
        :param error_msg: 错误信息
        :param task_id: 任务ID
        """
        try:
            # 如果之前有失败记录，发送恢复通知
            if mon_task_obj.get().num_of_fail != 0:
                mon_task_obj.update(recovery_time=timezone.now())
                self.__send_alert_msg(mon_task_obj, status=1, error_msg=error_msg, task_id=task_id)
            # 重置失败计数和状态
            mon_task_obj.update(task_status=1, num_of_fail=0)
        except Exception as e:
            logger.error(e)
            raise

    def __set_fail_num_and_send_fail_msg(self, mon_task_obj, error_msg='', task_id=0):
        """
        处理任务失败的情况
        :param mon_task_obj: MonitorTask查询集
        :param error_msg: 错误信息
        :param task_id: 任务ID
        """
        try:
            fail_num = mon_task_obj.get().num_of_fail  # 已失败次数
            if fail_num == 0:
                mon_task_obj.update(fail_time=timezone.now())  # 记录首次失败时间

            trigger_threshold = mon_task_obj.get().trigger_threshold  # 告警阈值

            # 达到告警阈值时处理告警
            if fail_num >= int(trigger_threshold):
                timestamp = time.time()
                timestamp_now = int(timestamp)

                # 检查沉默时间是否已过
                if timestamp_now - mon_task_obj.get().last_alert_time >= int(mon_task_obj.get().quiet_time) * 60:
                    mon_task_obj.update(last_alert_time=timestamp_now)
                    self.__send_alert_msg(mon_task_obj, status=0, error_msg=error_msg, task_id=task_id)

            # 增加失败计数并更新状态
            fail_num += 1
            mon_task_obj.update(task_status=2, num_of_fail=fail_num)
        except Exception as e:
            logger.error(e)
            raise

    def __send_alert_msg(self, obj, status=0, error_msg='', task_id=0):
        """
        发送告警/恢复消息
        :param obj: MonitorTask查询集
        :param status: 0-告警 1-恢复
        :param error_msg: 错误信息
        :param task_id: 任务ID
        """
        try:
            # 获取告警相关信息
            alert_type = obj.get().get_type_display()
            group = obj.values('mon_group__name').first()['mon_group__name']
            name = obj.get().name
            mon_path = obj.get().mon_path
            start_time = obj.get().fail_time
            end_time = obj.get().recovery_time
            host = obj.values('host__hostname').first()['host__hostname']

            # 获取告警联系人组信息
            alert_recv_group_obj = (AlertContactGroup.objects.filter(monitortaskalertgroup__monitor_task_id=task_id)
                                    .values('id', 'name'))
            group_id_list = [item['id'] for item in alert_recv_group_obj]

            # 获取联系人信息
            person_msg = AlertContactPerson.objects.filter(
                alertcontactpersongroup__group_id__in=group_id_list).values().distinct()

            # 拼接联系人组名称
            group_names = ""
            for index in range(len(alert_recv_group_obj)):
                if index == len(alert_recv_group_obj) - 1:
                    group_names += alert_recv_group_obj[index]['name']
                    break
                group_names += alert_recv_group_obj[index]['name'] + ','

            # 计算持续时间并创建告警日志
            timestamp = time.time()
            timestamp_now = int(timestamp)
            duration_time = int(math.ceil((timestamp_now - obj.get().fail_time.timestamp()) / 60))
            AlertLogs.objects.create(
                name_id=obj.get().id,
                mon_group=group,
                type=alert_type,
                mon_path=mon_path,
                alert_status=status,
                recv_alert_group=group_names,
                duration_time=duration_time,
                first_alert_time=obj.get().fail_time
            )

            # 1、获取person_msg中的每个联系人信息
            # 2、获取选取的所有告警方式
            # 3、判断监控类型，进程检测和自定义脚本中需要选择主机，其他的不需要
            for item in person_msg:
                for m in obj.get().alert_method_list:
                    if obj.get().type == 1 or obj.get().type == 2:
                        self.__send(obj, item, alert_type, group, name, mon_path, start_time, end_time, host, error_msg,
                                    status, 0, m)
                    if obj.get().type == 3 or obj.get().type == 4:
                        self.__send(obj, item, alert_type, group, name, mon_path, start_time, end_time, host, error_msg,
                                    status, 1, m)
        except Exception as e:
            logger.error(e)
            raise

    def __send(self, obj, args, alert_type, group, name, mon_path, start_time, end_time, hostname,
               error_msg, status, host, alert_method):
        """
        实际发送消息的方法
        :param obj: MonitorTask查询集
        :param args: 联系人信息
        :param alert_type: 告警类型
        :param group: 组名
        :param name: 任务名
        :param mon_path: 监控路径
        :param start_time: 开始时间
        :param end_time: 结束时间
        :param hostname: 主机名
        :param error_msg: 错误信息
        :param status: 状态
        :param host: 是否包含主机信息
        :param alert_method: 告警方式
        """
        # 构造告警信息字典
        alert_dict = {
            'type': alert_type,
            'group': group,
            'name': name,
            'obj': mon_path,
            'start_time': start_time,
            'end_time': end_time,
            'hostname': hostname,
            'error': error_msg,
            'status': status,
            'host': host,
        }

        try:
            # 根据告警方式调用不同的发送方法
            if alert_method == "1":  # 邮件
                send_email("告警测试", args['email'], alert_msg=1, alert_dict=alert_dict)
            if alert_method == "2":  # 飞书
                send_feishu_message(args['feishu_bot'], args['feishu_bot_sign'], alert_dict)
            if alert_method == "3":  # 钉钉
                send_dingtalk_message(args['ding_talk'], args['ding_talk_sign'], alert_dict)
            if alert_method == "4":  # 企业微信
                send_wxwork_message(
                    args['qw_robot_secret'],
                    args['qw_id'],
                    args['qw_agent_id'],
                    args['qw_to_user'],
                    alert_dict
                )
        except Exception as e:
            logger.error(e)
            raise
